import pandas as pd
from prefect import Task, Flow


class BaseTask(Task):
    def __init__(self, **kwargs):
        super().__init__(fn=self.run, **kwargs)  # 关键：把 run 传进去

    def run(self, df: pd.DataFrame) -> pd.DataFrame:
        """
           对输入的 DataFrame 进行结构化处理，生成信息相关的多张表。

           :param df: 输入的原始 DataFrame
           :return: 包含信息表、标签表、附件表、章节表的字典
           """
        raise NotImplementedError

    @staticmethod
    def clean_dataframe(dataframe: pd.DataFrame) -> pd.DataFrame:
        """
        将 DataFrame 中的 NaN 值替换为 None，以避免 MySQL 插入错误。
        """
        return dataframe.where(pd.notnull(dataframe), None)